package com.lx.cozeai.web;


import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.coze.openapi.client.chat.CreateChatReq;
import com.coze.openapi.client.chat.model.ChatEvent;
import com.coze.openapi.client.chat.model.ChatEventType;
import com.coze.openapi.client.connversations.message.model.Message;
import com.coze.openapi.service.auth.TokenAuth;
import com.coze.openapi.service.config.Consts;
import com.coze.openapi.service.service.CozeAPI;
import com.lx.cozeai.common.Result;
import com.lx.cozeai.dao.ChatMessageMapper;
import com.lx.cozeai.entity.ChatMessage;
import com.lx.cozeai.entity.User;
import io.reactivex.Flowable;
import io.reactivex.disposables.Disposable;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.io.IOException;
import java.time.LocalDateTime;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

@RestController
@RequestMapping("/chat")
public class ChatController {

    private static final Logger log = LogManager.getLogger(ChatController.class);
    // 使用线程安全的缓冲区
    private final ConcurrentMap<Long, ChatBuffer> chatBuffers = new ConcurrentHashMap<>();

    @Autowired
    ChatMessageMapper chatMessageMapper;

    @Value("${coze.token}")
    String token;
    @Value("${coze.bot_id}")
    String botID;
    @Value("${coze.user_id}")
    String userID;


    @GetMapping("/stream")
    public SseEmitter chatStream(@RequestParam String msg,
                                 HttpServletResponse response,
                                 HttpSession session) throws IOException {
        // 1. 初始化SSE连接
        SseEmitter emitter = new SseEmitter(180_000L); // 3分钟超时
        response.setHeader("Access-Control-Allow-Origin", "*");

        // 2. 获取用户信息
        User currentUser = (User) session.getAttribute("myinfo");
        Long uid = Long.valueOf(currentUser != null ? currentUser.getId() : null);

        try {
            // 3. 初始化Coze API客户端
            CozeAPI coze = initCozeClient();

            // 4. 构建聊天请求
            CreateChatReq req = buildChatRequest(msg, currentUser);

            // 5. 获取聊天事件流
            Flowable<ChatEvent> resp = coze.chat().stream(req);

            // 6. 为当前用户初始化缓冲区
            if (currentUser != null) {
                chatBuffers.putIfAbsent( uid, new ChatBuffer(msg));
            }

            // 7. 订阅事件流
            Disposable subscription = resp.subscribe(
                    event -> handleStreamEvent(event, emitter, uid),
                    error -> handleStreamError(error, emitter,  uid),
                    () -> completeStream(emitter, uid)
            );

            // 8. 设置Emitter生命周期回调
            emitter.onCompletion(() -> cleanupResources(subscription, uid));
            emitter.onTimeout(() -> handleTimeout(emitter, subscription, uid));

        } catch (Exception e) {
            log.error("Initialization error", e);
            emitter.completeWithError(e);
            cleanupResources(null, uid); // 清理资源
        }

        return emitter;
    }

    // === 辅助方法 ===

    private CozeAPI initCozeClient() {
        TokenAuth authCli = new TokenAuth(token);
        return new CozeAPI.Builder()
                .baseURL(Consts.COZE_CN_BASE_URL)
                .auth(authCli)
                .readTimeout(80000)
                .build();
    }

    private CreateChatReq buildChatRequest(String msg, User user) {

        return CreateChatReq.builder()
                .botID(botID)
                .userID(userID)
                .messages(Collections.singletonList(Message.buildUserQuestionText(msg)))
                .build();
    }

    private void handleStreamEvent(ChatEvent event, SseEmitter emitter, Long uid) {
        if (!ChatEventType.CONVERSATION_MESSAGE_DELTA.equals(event.getEvent())) {
            return;
        }

        String content = event.getMessage().getContent();
        if (content == null || content.trim().isEmpty()) {
            return;
        }

        try {
            // 1. 实时发送数据到前端
            emitter.send(SseEmitter.event()
                    .data(content)
                    .id(UUID.randomUUID().toString())
                    .comment("message chunk"));

            // 2. 缓冲AI响应内容
            if (uid != null) {
                ChatBuffer buffer = chatBuffers.get(uid);
                if (buffer != null) {
                    buffer.appendAiResponse(content);
                }
            }
        } catch (IOException e) {
            log.error("Failed to send SSE event", e);
            emitter.completeWithError(e);
        }
    }

    private void handleStreamError(Throwable error, SseEmitter emitter, Long uid) {
        log.error("Chat stream error", error);
        emitter.completeWithError(error);
        saveConversation(uid); // 尝试保存已接收的部分
    }

    private void completeStream(SseEmitter emitter, Long uid) {
        emitter.complete();
        saveConversation(uid); // 保存完整对话
    }

    private void cleanupResources(Disposable subscription, Long uid) {
        if (subscription != null && !subscription.isDisposed()) {
            subscription.dispose();
        }
        if (uid != null) {
            chatBuffers.remove(uid); // 清理缓冲区
        }
    }

    private void handleTimeout(SseEmitter emitter, Disposable subscription, Long uid) {
        log.warn("SSE connection timed out for user: {}", uid);
        emitter.complete();
        saveConversation(uid); // 超时也尝试保存
        cleanupResources(subscription, uid);
    }

    private synchronized void saveConversation(Long uid) {
        if (uid == null) return;

        ChatBuffer buffer = chatBuffers.get(uid);
        if (buffer == null) return;

        try {
            // 保存用户消息
            if (!buffer.getUserMessage().isEmpty()) {
                ChatMessage userMsg = new ChatMessage();
                userMsg.setContent(buffer.getUserMessage());
                userMsg.setUid(Math.toIntExact(uid));
                userMsg.setTime(LocalDateTime.now());
                userMsg.setType(1);
                chatMessageMapper.insert(userMsg);
            }

            // 保存AI响应
            String aiResponse = buffer.getAiResponse();
            if (!aiResponse.isEmpty()) {
                ChatMessage aiMsg = new ChatMessage();
                aiMsg.setContent(aiResponse);
                aiMsg.setUid(Math.toIntExact(uid));
                aiMsg.setTime(LocalDateTime.now());
                aiMsg.setType(0);
                chatMessageMapper.insert(aiMsg);
            }

            log.debug("Conversation saved for user: {}", uid);
        } catch (Exception e) {
            log.error("Failed to save conversation for user: {}", uid, e);
        } finally {
            buffer.clear(); // 清空缓冲区
        }
    }

    /**
     * 内部类：聊天缓冲区
     */
    private static class ChatBuffer {
        private final String userMessage;
        private final StringBuilder aiResponse = new StringBuilder();

        ChatBuffer(String userMessage) {
            this.userMessage = userMessage;
        }

        void appendAiResponse(String content) {
            aiResponse.append(content);
        }

        String getUserMessage() {
            return userMessage;
        }

        String getAiResponse() {
            return aiResponse.toString();
        }

        void clear() {
            aiResponse.setLength(0);
        }
    }



    @RequestMapping("/getMessage")
    public Result getMessage(int uid) {
        LambdaQueryWrapper<ChatMessage> lqw = new LambdaQueryWrapper<>();
        lqw.eq(ChatMessage::getUid, uid);
        List<ChatMessage> chatMessages = chatMessageMapper.selectList(lqw);
        return new Result(1, "查询成功", chatMessages);
    }

}
